#!/usr/bin/env python

# a stupid validator for Lys e-groups in a pdb file
# given a silk result
# please don't comment on this code...

import sys, xlh, getopt
from math import pow, sqrt
from numpy import base_repr


aa3to1 = {
'Mse': 'X',
'Ala': 'A',
'Cys': 'C',
'Asp': 'D',
'Glu': 'E',
'Phe': 'F',
'Gly': 'G',
'His': 'H',
'Ile': 'I',
'Lys': 'K',
'Leu': 'L',
'Met': 'M',
'Asn': 'N',
'Pro': 'P',
'Gln': 'Q',
'Arg': 'R',
'Ser': 'S',
'Thr': 'T',
'Val': 'V',
'Trp': 'W',
'Tyr': 'Y',
'Hyp': 'P',
'Ilx': 'I',
'Trx': 'Y',
'Csx': 'C' 
}

LysArray = []
CoordArray = []
Sequences = []	# original protein sequences
ToRemove = {} # a dictionary chain->[positions] to store aa to remove (i.e. leading peptides)
							# to keep positions in frame with the sequences
fixmod1 = 'carbamidomethyl C'
xlinkLight = 'x-BS2GD0 (K)'
xlinkHeavy = 'x-BS2GD4 (K)'
Hydrogen = xlh.Elements.getfloat("H", "mmass")
PdbFile = ''
SilkFile = ''
MappingFile = ''
oldId = ''
seq = ''
MinDistance = 5
MaxDistance = 35
Tolerance = 10
InteractionList =xlh.XLinkList()
StructMap = {}

CheckMgf = False	# True if mgf file are provided


def distance(a, b):
	# a and b are two lists
	s = 0.0
	for x, y in zip(a, b):
		s += pow(x-y, 2)
	
	return sqrt(s)

def isWithin(Peptide, position):
	starts = Peptide.startPositions()
	ends = Peptide.endPositions()
	
	for x in range(len(starts)):
		if position > starts[x] and position < ends[x]:
			return True
			
	return False

def mr2mz(mass, charge):
	return (mass + charge * Hydrogen) / charge
	
def filterPeptides(Protein, position):
	Peptides = filter(lambda p: isWithin(p, position), Protein.peptides())
	Peptides = filter(lambda p: p.length() >= 3, Peptides)
	Peptides = filter(lambda p: 'K' in p.sequence()[:-1] or p.isNTerminal(), Peptides)
	return Peptides

def finneganScan(descr):
#	p = descr.find('precIntensity')
	p = descr.find('FinneganScanNumber')
	return descr[p:]


def checkDoublet(spectrum, Tolerance):
	return True
	intFilter = 3	# minimum intensity for a peak
	Doublet = 2		# mass distance between doublet peaks
	intRatio = 0.7	# intensity ratio between doublet peaks
	
	spsl = spectrum.spectrumSlice(min = spectrum.precursorMz()).filterOnIntensity(intFilter)
#	print spectrum.length()
	for peak in spsl.items():
		rtol = peak[0] / 1000000 * Tolerance

		clpeaks = filter(lambda x: x[0] - peak[0] <= (Doublet + rtol) and x[0] - peak[0] >= (Doublet - rtol), spectrum.peaks()) 

		if filter(lambda x: x[1] / peak[1] >= intRatio or peak[1] / x[1] <= 1 / intRatio, clpeaks):
			#return spsl.closeTo(peak[0], (Doublet + rtol)).peaks()
			return True
			
		# if a doublet in 2 Da has been found stop here, otherwise check if there's a doublet in 4 Da...
			
		clpeaks = filter(lambda x: x[0] - peak[0] <= (2*Doublet + rtol) and x[0] - peak[0] >= (2*Doublet - rtol), spectrum.peaks()) 

		if filter(lambda x: x[1] / peak[1] >= intRatio or peak[1] / x[1] <= 1 / intRatio, clpeaks):
			#return spsl.closeTo(peak[0], (2 * Doublet + rtol)).peaks()
			return True
			
	return False

def exitOnError(errmsg, exit):
	print 'The following error(s) occurred:\n'
	print errmsg
	sys.exit(exit)
def usage():
	print """Usage:

LysCheck.py [-h] -p PdbFile -f SilkResults -c ChainMapppings[-t TOLERANCE] [-m Min] [-M Max] 

 -h:	Print this help and exit
 -p:	PDB File to check
 -f:	Silk results file
 -c:	Chain Mapping file
 -t:	Tolerance for MS comparison (ppm) [Default 10]
 -m:	Minimum distance between Lysin residues [Default 5]
 -M:	Maximum distance between Lysin residues [Default 25, 0 for infinite]
	"""
#	sys.exit(1)

try:
	optlist, args = getopt.getopt(sys.argv[1:], 'hp:f:c:t:m:M:')
except getopt.GetoptError:
	usage()

if not optlist:
	usage()

for o, a in optlist:
	if o == '-h':
		usage()
	if o == '-p':
		PdbFile = a
	if o == '-t':
		Tolerance = int(a)
	if o == '-m':
		MinDistance = int(a)
	if o == '-M':
		MaxDistance = int(a)
		if MaxDistance == 0:
			MaxDistance = 1000000
	if o == '-f':
		SilkFile = a
	if o == '-c':
		MappingFile = a

if not PdbFile:
	exitOnError("Missing PDB file", 1)
#if not SilkFile:
#	exitOnError("Missing Silk file", 1)
if not MappingFile:
	exitOnError("Missing chain mapping file", 1)
	
ChainMappings = {}
ChainCorrections = {}
for line in open(MappingFile):
	# associate an accession to each chain
	l = line.strip().split()
	ChainMappings[l[0]] = l[1]
	try:
		ChainCorrections[l[1]] = int(l[2])
	except IndexError:
		ChainCorrections[l[1]] = 0
		
###
# Parse Silk results to get all events
###
if SilkFile:
	eLengths = []
	ePerc = []
	for line in open(SilkFile, 'r'):
		if '[ Interaction' in line:
			tmp = xlh.Interaction()
		if 'Coordinates' in line:
			(left, right) = line[13:].strip().split()
			(pr, pos) = left.split(':')
			xpr = xlh.Protein()
			xpr.setAccession(pr)
			tmp.setNProtein(xpr)
			tmp.setStart(pos)
			
			(pr, pos) = right.split(':')
			xpr = xlh.Protein()
			xpr.setAccession(pr)
			tmp.setCProtein(xpr)
			tmp.setEnd(pos)
		if 'G-Score' in line:
			tmp._gScore = float(line.split()[1])
			tmp._calculatedGScore = True
			ePerc.append(line.split()[3])
		if 'Rank' in line:
			tmp._rank = int(line.split()[1])
			InteractionList.addInteraction(tmp)
		if 'N. of spectra:' in line:
			eLengths.append(line.split()[3])
		if '#######' in line:
			break
###


f = open(PdbFile, 'r')
for line in f:
	if line[:6] == 'SEQRES':
		#sequence information
		chainId = line[11]
		tmp = ''.join(map(lambda x: aa3to1[x.capitalize()], line[19:].split()))
		if chainId == oldId:
			ThisProt.addSequence(tmp)
		else:
			ThisProt = xlh.Protein()
			ThisProt.setSequence(tmp)
			ThisProt.setAccession(chainId)
			Sequences.append(ThisProt)
			oldId = chainId	
		
	if line[:6] == 'SEQADV':
		# collect information about missing aminoacids, we will remove them from
		# sequences. Actually, we are interested in removing aa with negative position
		# as this means that a leading peptide has been somehow cleaved or eliminated
		# from sequence. Other aa missing... well, don't consider them by now
		chainId = line[16]
		position = int(line[18:22])
		if position <= 0:
			try:
				ToRemove[chainId].append(position)
			except KeyError:
				ToRemove[chainId] = [position]
	
	if line[:4] == 'ATOM' or line[:4] == 'TER ':
	
		chainId = line[21]
		aaPos = int(line[22:26])
		aa = line[17:20]
		atomId = line[12:16]

		if atomId == ' NZ ' and aa == 'LYS':
			StructMap[line[21:26]] = map(lambda x: float(x), [line[30:38], line[38:46], line[46:54]])
			LysArray.append([chainId,line[22:26]])
			CoordArray.append(map(lambda x: float(x), [line[30:38], line[38:46], line[46:54]]))
		if atomId == ' N  ' and aaPos == 1:
			StructMap[line[21:26]] = map(lambda x: float(x), [line[30:38], line[38:46], line[46:54]])
			LysArray.append([chainId,line[22:26]])
			CoordArray.append(map(lambda x: float(x), [line[30:38], line[38:46], line[46:54]]))
			
f.close()

# as we now process aa in ToRemove, Sequences and XtalSeq entries should start with the 
# same aminoacid, that means we do not need XtalSeq anymore

for id in ToRemove.keys():
	protein = filter(lambda x: x.accession() == id, Sequences)[0]
	# assume that the list of aa to remove is a range from -x to 0
	# instead of collecting position we could just count the missing leading aa
	# but, you never know with pdb, so just keep the list and maybe thnk to a 
	# smartest way to do this....
	protein.setSequence(protein.sequence()[-(len(protein.sequence()) - len(ToRemove[id])):])
	

# cleave each protein
map(lambda x: x.cleave(), Sequences)

#		print "Distance\tChain \tPeptide 1\tHy 1\tStart\tK pos\tEnd\tChain 2\tPeptide 2\tHy 2\tStart\tK pos\tEnd\tLight Mr\tHeavy Mr"

if SilkFile:
	x = 1
	print "No,Match,Prot1,start,Prot2,end,distance,gScore,rank,nSpectra,sPerc,HeadTail,MultiSeq,MultiXL,BothPos,MultiSpec"
	
	for x in range(InteractionList.eLength()):
		interaction = InteractionList._eList[x]
		# build N accession
	
		try:
			nChain = ChainMappings[interaction.NProtein().accession()]
		except KeyError:
			sys.stderr.write(str(interaction.NProtein().accession() + " not in PDB file\n"))
			continue
		try:
			rChain = ChainMappings[interaction.CProtein().accession()]
		except KeyError:
			sys.stderr.write(str(interaction.CProtein().accession() + " not in PDB file\n"))
			continue
			
	
		nid = nChain + str(interaction.start() - ChainCorrections[nChain]).rjust(4)
		rid = rChain + str(interaction.end() - ChainCorrections[rChain]).rjust(4)
		Coord1 = []
		Coord2 = []
		try:
			Coord1 = StructMap[nid]
		except KeyError:
			continue
		try:
			Coord2 = StructMap[rid]
		except KeyError:
			continue
		TheDistance = distance(Coord1, Coord2)
		TVal = 'F'
		if TheDistance >= MinDistance and TheDistance <= MaxDistance:
			TVal = 'T'
		outList = [str(1+x), TVal, interaction.NProtein().accession(), str(interaction.start()),
		interaction.CProtein().accession(), str(interaction.end()), str(TheDistance), str(interaction.gScore()), str(interaction.rank()), str(eLengths[x]),ePerc[x]] +	map(lambda l:l, base_repr(interaction.rank(), 2).rjust(5, '0'))
		
		outString = ','.join(outList)
		print outString
		x += 1
else:
	chainIDs = StructMap.keys()
	chainIDs.sort()
	reversedMappings = {}
	for (x, y) in ChainMappings.items():
		reversedMappings[y] = x
#		print reversedMappings
	print "Prot1,start,Prot2,end,distance"
	for x in range(len(chainIDs)):
		nid = chainIDs[x]
		for y in range(x, len(chainIDs)):
			rid = chainIDs[y]
			nChain = nid[0]
			try:
				start = int(nid[1:]) + ChainCorrections[nChain]
			except KeyError:
				continue
			rChain = rid[0]
			try:
				end = int(rid[1:]) + ChainCorrections[rChain]
			except KeyError:
				continue
			TheDistance = distance(StructMap[nid], StructMap[rid])
			if TheDistance >= MinDistance and TheDistance <= MaxDistance:
				try:
					outList = [reversedMappings[nChain], str(start), reversedMappings[rChain], str(end), str(TheDistance)]
					print ','.join(outList)
				except KeyError:
					continue
		

sys.exit(0)
